Skip to content

[RunInference] Raise when a model handler returns fewer predictions than inputs - #40929

Open
tannnmayy wants to merge 1 commit into
apache:masterfrom
tannnmayy:strict-result-pairing
Open

tannnmayy wants to merge 1 commit into
apache:masterfrom
tannnmayy:strict-result-pairing

Conversation

@tannnmayy

@tannnmayy tannnmayy commented Oct 31, 2026 •

Copy link
Copy Markdown
Contributor

Fixes #40928

utils._convert_to_result pairs each batch of inputs with the model's predictions using zip(), which stops at the shorter of the two. When a handler produced fewer predictions than inputs, the leftover inputs got no result and no error, and num_inferences still counted them. That's how #36587 and #40913 went unnoticed, and it still happens with built-in handlers, for example with PyTorch or Keras models that return a tuple of outputs, and with Hugging Face causal LMs that return past_key_values (details in the issue).

This PR:

  • makes _convert_to_result raise a ValueError when there are fewer predictions than inputs. The message gives the counts, names any dict output that doesn't have one entry per input, and suggests a fix (an inference function that returns one prediction per input, or inference_args={'use_cache': False} for Hugging Face models);
  • leaves extra predictions ignored, as before, so pipelines in which every input already gets a result keep working (for example, tuple outputs with batch size 1);
  • fixes three existing tests that only passed because of the truncation. In sklearn_inference_test.py, test_custom_inference_fn used an inference function that returned one prediction for three inputs, and its zip()-based check compared only the first result. In tensorflow_inference_test.py, the two env-var tests mocked the model, so their pipelines produced no results at all. The fakes now return one prediction per input, and the scikit-learn test also checks the number of results;
  • adds unit tests for the helper, a RunInference test that checks the affected inputs reach the failure output of with_exception_handling(), and a CHANGES.md entry.

Testing:

  • 9 new tests in utils_test.py. The 4 that check the new error fail without this change; the other 5 cover behavior that doesn't change.
  • The whole apache_beam/ml test suite with torch, TensorFlow, transformers, ONNX, scikit-learn and XGBoost installed (Python 3.12): 609 passed, against 600 on master. The only differences are the 9 new tests. The tests that can't run locally (they need GCP credentials or Docker) fail the same way on both.
  • The real-model cases from the issue: every case that dropped inputs now raises, and the single-output controls still return one result per input.
  • ruff, yapf and pyrefly report no issues.

The issue also asks whether RunInference should check the count itself, to cover custom handlers that pair results with their own zip(). I left that out to keep this change to the shared helper. If you'd prefer a warning instead of an error, or the check in RunInference itself, I'm happy to change it.

R: @jrmccluskey

…han inputs

utils._convert_to_result paired inputs and predictions with zip(), which
stops at the shorter of the two, so inputs without a prediction were
dropped without an error. It now raises a ValueError that gives the counts
and names any dict output without one entry per input. Extra predictions
are still ignored, so pipelines in which every input already gets a result
are unaffected.

Also fixes three tests that only passed because of the truncation.

Fixes apache#40928
@github-actions github-actions Bot added python size/m M: 50-249 lines changed labels Oct 31, 2026

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📋 Review Summary

This PR updates utils._convert_to_result to raise a descriptive ValueError when a model handler produces fewer predictions than inputs in a batch, preventing silent data loss caused by zip() truncation. It also updates three existing unit tests that inadvertently relied on this truncation behavior and adds comprehensive test coverage for the new validation.

🔍 General Feedback

  • The implementation in utils.py cleanly handles both sequence/generator and dictionary prediction outputs, providing clear, actionable error messages (including identifying specific short keys in dictionary outputs).
  • Keeping extra predictions ignored preserves backward compatibility for existing single-element batch pipelines while catching silent data drops.
  • The test additions in utils_test.py and fixes in sklearn_inference_test.py and tensorflow_inference_test.py are thorough and well-structured.

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

@codecov

codecov Bot commented Oct 31, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 89.65517% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.09%. Comparing base (12bbb17) to head (06d1dc8).
⚠️ Report is 19 commits behind head on master.

Files with missing lines Patch % Lines
sdks/python/apache_beam/ml/inference/utils.py 89.65% 3 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##           master   #40929      +/-   ##
==========================================
+ Coverage   79.83%   80.09%   +0.26%     
==========================================
  Files         494      494              
  Lines       83575    83628      +53     
==========================================
+ Hits        66722    66985     +263     
+ Misses      16853    16643     -210     
Flag Coverage Δ
python 80.09% <89.65%> (+0.26%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@github-actions

Copy link
Copy Markdown
Contributor

Assigning reviewers:

R: @claudevdm for label python.
Backup reviewer: shunping (if the main reviewer is unresponsive, you can escalate to them by tagging their username)

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

  • stop reviewer notifications - opt out of the automated review tooling
  • remind me after tests pass - tag the comment author after tests pass
  • waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]: RunInference silently drops inputs when a model returns fewer predictions than inputs

1 participant